﻿using System.Collections;
using System.Collections.Generic;
using UnityEngine;

public class GraphNeuralNetwork : MonoBehaviour
{
    public GameObject graphNeuralPrefab;
    public GameObject graphWeightPrefab;
    public List<GraphNeuralLayerData> graphNeuralLayerList = new List<GraphNeuralLayerData>();

    public void CreateGraphNeuralNetwork(NeuralNetwork nn)
    {
        float x = 0;
        float z = 0;
        for (int i = 0; i < nn.neuralLayers.Count; i++)
        {

            graphNeuralLayerList.Add(new GraphNeuralLayerData());

            x = transform.position.x + i * 6;

            for (int j = 0; j < nn.neuralLayers[i].neurals.Count; j++)
            {

                z = transform.position.z + j * 3;
                GameObject go = Instantiate(graphNeuralPrefab) as GameObject;
                go.transform.position = new Vector3(x, 20, z);
                GraphNeural un = go.GetComponent<GraphNeural>();
                un.neuralIndex = j;
                un.layerIndex = i;
                go.transform.SetParent(transform);

                graphNeuralLayerList[i].GraphNeuralList.Add(new GraphNeuralData()
                {
                    isBias = false,
                    layerIndex = i,
                    neuralIndex = j,
                    neuralObj = go
                });

                if (i == 0)
                {
                    GameObject line = Instantiate(graphWeightPrefab) as GameObject;
                    line.GetComponent<LineRenderer>().SetPosition(0, go.transform.position);
                    line.GetComponent<LineRenderer>().SetPosition(1, go.transform.position - go.transform.right * 3);
                    line.transform.SetParent(transform);

                    graphNeuralLayerList[i].GraphNeuralList[j].graphWeightList.Add(new GraphWeightData()
                    {
                        fromLayerIndex = -2,
                        fromNeuralIndex = -2,
                        weight = 1,
                        weightObj = line
                    });
                }

                for (int k = 0; k < nn.neuralLayers[i].neurals[j].weights.Length; k++)
                {
                    GameObject line = Instantiate(graphWeightPrefab) as GameObject;
                    line.GetComponent<LineRenderer>().SetPosition(0, go.transform.position);
                    line.GetComponent<LineRenderer>().SetPosition(1, new Vector3((i - 1) * 6, 20, -k * 3));
                    line.transform.SetParent(transform);

                    graphNeuralLayerList[i].GraphNeuralList[j].graphWeightList.Add(new GraphWeightData()
                    {
                        fromLayerIndex = i - 1,
                        fromNeuralIndex = k == nn.neuralLayers[i].neurals[j].weights.Length - 1 ? -1 : k,
                        weight = 1,
                        weightObj = line
                    });
                }
            }

            if (i != nn.neuralLayers.Count - 1)
            {
                GameObject go2 = Instantiate(graphNeuralPrefab) as GameObject;
                go2.transform.position = new Vector3(x, 20, transform.position.z + nn.neuralLayers[i].neurals.Count * 3);
                GraphNeural un2 = go2.GetComponent<GraphNeural>();
                un2.neuralIndex = nn.neuralLayers[i].neurals.Count;
                un2.layerIndex = i;
                go2.transform.SetParent(transform);

                graphNeuralLayerList[i].GraphNeuralList.Add(new GraphNeuralData()
                {
                    isBias = true,
                    layerIndex = i,
                    neuralIndex = -1,
                    neuralObj = go2
                });
            }
        }

        //transform.localRotation = Quaternion.Euler(0, 90, 0);
    }

    public void SetGraphNeuralNetwork(NeuralNetwork nn)
    {
        if (nn == null)
        {
            return;
        }
        for (int i = 0; i < graphNeuralLayerList.Count; i++)
        {
            for (int j = 0; j < graphNeuralLayerList[i].GraphNeuralList.Count; j++)
            {
                GraphNeural un = graphNeuralLayerList[i].GraphNeuralList[j].neuralObj.GetComponent<GraphNeural>();
                un.layerIndex = i;
                un.neuralIndex = j;
                if (i != graphNeuralLayerList.Count - 1 && j == graphNeuralLayerList[i].GraphNeuralList.Count - 1)
                {
                    un.nn = nn;
                    un.isBias = true;
                }
                else
                {
                    un.nn = nn;
                    un.isBias = false;


                    for (int k = 0; k < graphNeuralLayerList[i].GraphNeuralList[j].graphWeightList.Count; k++)
                    {
                        GraphWeight ul = graphNeuralLayerList[i].GraphNeuralList[j].graphWeightList[k].weightObj.GetComponent<GraphWeight>();
                        ul.fromLayerIndex = graphNeuralLayerList[i].GraphNeuralList[j].graphWeightList[k].fromLayerIndex;
                        ul.fromNeuralIndex = graphNeuralLayerList[i].GraphNeuralList[j].graphWeightList[k].fromNeuralIndex;
                        ul.nn = nn;
                        if (ul.fromNeuralIndex != -2)
                        {
                            ul.weight = nn.neuralLayers[i].neurals[j].weights[k];
                            try
                            {
                                if (ul.fromNeuralIndex == -1)
                                {
                                    ul.startPos = graphNeuralLayerList[ul.fromLayerIndex].GraphNeuralList[graphNeuralLayerList[ul.fromLayerIndex].GraphNeuralList.Count - 1].neuralObj.transform;
                                }
                                else
                                {
                                    ul.startPos = graphNeuralLayerList[ul.fromLayerIndex].GraphNeuralList[ul.fromNeuralIndex].neuralObj.transform;
                                }
                            }
                            catch
                            {
                                Debug.LogError(ul.fromLayerIndex + "," + ul.fromNeuralIndex);
                            }
                            ul.endPos = un.transform;
                        }
                        else
                        {
                            ul.weight = 1;
                            ul.startPos = null;
                            ul.endPos = un.transform;
                        }

                    }
                }

            }
        }
    }
}

public class GraphNeuralData
{
    public GameObject neuralObj;
    public int layerIndex;
    public int neuralIndex;
    public bool isBias;

    public List<GraphWeightData> graphWeightList = new List<GraphWeightData>();
}

public class GraphNeuralLayerData
{
    public List<GraphNeuralData> GraphNeuralList = new List<GraphNeuralData>();
}

public class GraphWeightData
{
    public GameObject weightObj;
    public int fromLayerIndex;
    public int fromNeuralIndex;
    public double weight;
}
